import {
  Tensor,
  concat2d,
  slice,
  add,
  div,
  sub,
  mul,
  tidy,
  squeeze,
  concat,
  tensor1d,
} from "@tensorflow/tfjs-core";
import { scaleBox } from "./box";
import type { Tensor1D, Tensor2D } from "@tensorflow/tfjs-core";
import type { FacePrediction, NormalizedFace } from "./face";
import type { Box } from "./box";

export const generateAnchors = (width: number, height: number) => {
  const outputSpec = {
    strides: [8, 16],
    anchors: [2, 6],
  };
  const anchors = [];
  for (let i = 0; i < outputSpec.strides.length; i++) {
    const stride = outputSpec.strides[i];
    const gridRows = Math.floor((height + stride - 1) / stride);
    const gridCols = Math.floor((width + stride - 1) / stride);
    const anchorsNum = outputSpec.anchors[i];

    for (let gridY = 0; gridY < gridRows; gridY++) {
      const anchorY = stride * (gridY + 0.5);

      for (let gridX = 0; gridX < gridCols; gridX++) {
        const anchorX = stride * (gridX + 0.5);
        for (let n = 0; n < anchorsNum; n++) {
          anchors.push([anchorX, anchorY]);
        }
      }
    }
  }

  return anchors;
};

export const decodeBounds = (
  boxOutputs: Tensor2D,
  anchors: Tensor2D,
  inputSize: Tensor1D,
): Tensor2D => {
  const boxStarts = slice(boxOutputs, [0, 1], [-1, 2]);
  const centers = add(boxStarts, anchors);
  const boxSizes = slice(boxOutputs, [0, 3], [-1, 2]);

  const boxSizesNormalized = div(boxSizes, inputSize);
  const centersNormalized = div(centers, inputSize);

  const halfBoxSize = div(boxSizesNormalized, 2);
  const starts = sub(centersNormalized, halfBoxSize);
  const ends = add(centersNormalized, halfBoxSize);

  const startNormalized = mul(starts, inputSize);
  const endNormalized = mul(ends, inputSize);

  const concatAxis = 1;
  return concat2d(
    [startNormalized as Tensor2D, endNormalized as Tensor2D],
    concatAxis,
  );
};

export const scaleBoxFromPrediction = (
  face: FacePrediction | Box,
  scaleFactor: Tensor1D | [number, number],
) => {
  return tidy(() => {
    let box;
    if (face.hasOwnProperty("box")) {
      box = (face as FacePrediction).box;
    } else {
      box = face;
    }
    return squeeze(scaleBox(box as Box, scaleFactor).startEndTensor);
  });
};

export const flipFaceHorizontal = (
  face: NormalizedFace,
  imageWidth: number,
): NormalizedFace => {
  let flippedTopLeft: [number, number] | Tensor1D,
    flippedBottomRight: [number, number] | Tensor1D,
    flippedLandmarks;

  if (face.topLeft instanceof Tensor && face.bottomRight instanceof Tensor) {
    const [topLeft, bottomRight] = tidy(() => {
      return [
        concat([
          slice(sub(imageWidth - 1, face.topLeft as Tensor), 0, 1),
          slice(face.topLeft as Tensor, 1, 1),
        ]) as Tensor1D,
        concat([
          sub(imageWidth - 1, slice(face.bottomRight as Tensor, 0, 1)),
          slice(face.bottomRight as Tensor, 1, 1),
        ]) as Tensor1D,
      ];
    });

    flippedTopLeft = topLeft;
    flippedBottomRight = bottomRight;

    if (face.landmarks != null) {
      flippedLandmarks = tidy(() => {
        const a: Tensor2D = sub(tensor1d([imageWidth - 1, 0]), face.landmarks!);
        const b = tensor1d([1, -1]);
        const product: Tensor2D = mul(a, b);
        return product;
      });
    }
  } else {
    const [topLeftX, topLeftY] = face.topLeft as [number, number];
    const [bottomRightX, bottomRightY] = face.bottomRight as [number, number];

    flippedTopLeft = [imageWidth - 1 - topLeftX, topLeftY];
    flippedBottomRight = [imageWidth - 1 - bottomRightX, bottomRightY];

    if (face.landmarks != null) {
      flippedLandmarks = (face.landmarks as number[][]).map((coord) => [
        imageWidth - 1 - coord[0],
        coord[1],
      ]);
    }
  }

  const flippedFace: NormalizedFace = {
    topLeft: flippedTopLeft,
    bottomRight: flippedBottomRight,
  };

  if (flippedLandmarks != null) {
    flippedFace.landmarks = flippedLandmarks;
  }

  if (face.probability != null) {
    flippedFace.probability =
      face.probability instanceof Tensor
        ? face.probability.clone()
        : face.probability;
  }

  return flippedFace;
};
